#!/usr/bin/env python

"""
    This script fits the residuals of PTC curves and use the fitted line to determine the following nonlinearities:
    1) 14-56 ke-
    2) 7-70 ke-

    usage: 
        residual_linear_fit_dir.py data nl_fit 850

    2025-05-22： 往信号强端增加一个拟合数据点
"""

import sys
from glob import glob
import os
from astropy.table import Table
import numpy as np
from scipy import odr
from scipy.stats import sigmaclip
import matplotlib.pyplot as plt
from scipy.ndimage import median_filter

cat_dir = sys.argv[1]
out_dir = sys.argv[2]
wave = sys.argv[3]

def char2bool(data):
    out = []
    for i in range(len(data)):
        if data[i] == 'True':
            out.append(True)
        else:
            out.append(False)
    return np.array(out)
    
def d2line(x,y,k,b):
    return abs(k*x-y+b)/(1+k**2)**0.5
    
def ptc_fit(tab_file, fig_outdir):
    tab = Table.read(tab_file,format='ipac')
    t = tab['t']
    x = tab['dn_bin']
    y = tab['var_bin']
    
    #calculate image fullwell
    peakx = np.argmax(x)
    full_image_dn = x[peakx]

    #calculate turning point
    dn = median_filter(x, 3)
    peak_dn = np.argmax(dn)
    peak_var = np.argmax(y)
    peak = min(peak_dn, peak_var)
    tp = x[peak]

    #select data for fitting
    valid0 = char2bool(tab['valid'].data)
    valid = valid0.copy()
    ptc10 = char2bool(tab['ptc10'].data)
    ptc20 = char2bool(tab['ptc20'].data)
    
    #idx0 = np.where(valid)[0][0]
    #idxh = np.argmax(y)
    #valid[idx0:idxh+1] = True
    #valid[idxh+1:] = False
    
    idx0 = np.where(valid)[0][0]
    #idxh = len(np.where(valid)[0])+1    # 往信号强端增加一个拟合数据点
    idxh = len(np.where(valid)[0])+2    # 往信号强端增加2个拟合数据点
    valid[idx0:idxh+1] = True
    valid[idxh+1:] = False
    
    xfit = x[valid]
    yfit = y[valid]
    
    #fitting valid points with linear model applying weights
    mod_l = odr.polynomial(1)
    data = odr.Data(xfit,yfit, wd=1 / np.power(xfit, 2), we=1 / np.power(yfit, 2))
    s0 = np.median(yfit/xfit)
    fitter_l = odr.ODR(data, mod_l, ifixb=[0, 1], beta0 = [0, s0])
    res_l = fitter_l.run()
    gain = 1/res_l.beta[1]
    y1_l = x/gain
    tp = tp*gain
    full_image = full_image_dn * gain

    #plot points and non-linearity
    plt.figure(figsize=[18,8])
    plt.subplot(121)
    plt.scatter(x/1000, y/1000, s=20, color = 'orange')
    plt.scatter(xfit/1000, yfit/1000, s=20, color = 'green')
    plt.xlabel("Signal (k ADU)")
    plt.ylabel("Corrected Variance (k ADU^2)")
    plt.plot(x/1000, x/gain/1000, color='red', linestyle='--', label='gain = '+str(np.round(gain,2))+'e-/ADU')
    plt.legend()
    
    #fitting non-linearity with linear model
    nl = (y*gain/x-1)*100
    mod_nl = odr.polynomial(1)
    data_nl = odr.Data(xfit, nl[valid])
    s_nl0 = np.median((nl[valid][1:]-nl[valid][0])/(xfit[1:]-xfit[0]))
    i_nl0 = nl[valid][0]
    fitter_nl = odr.ODR(data_nl, mod_nl, ifixb=[1, 1], beta0 = [i_nl0, s_nl0])
    res_nl = fitter_nl.run()
    i_nl = res_nl.beta[0]
    s_nl = res_nl.beta[1]
    ds = d2line(xfit,nl[valid],s_nl,i_nl)
    full = (-5-i_nl)/s_nl*gain
    if full > np.max(xfit) * gain:
        full = np.max(xfit)*gain

    #clipp outliers
    valid1 = np.array([False]*len(x))
    ds_clipped,_,_ = sigmaclip(ds,low=3.0,high=3.0)
    idx = []
    for i in range(len(ds_clipped)):
        if ds_clipped[i] in ds:
            idx.append(np.where(ds==ds_clipped[i])[0])
    idx = np.array(idx).astype('int')
    valid1[idx] = True

    #fit again
    data_nl_clipped = odr.Data(x[valid1], nl[valid1])
    s_nl0_clipped = np.median(nl[valid1]/x[valid1])
    i_nl0_clipped = nl[valid1][0]
    mod_nl_clipped = odr.polynomial(1)
    fitter_nl_clipped = odr.ODR(data_nl_clipped, mod_nl_clipped, ifixb=[1, 1], beta0=[i_nl0_clipped,s_nl0_clipped])
    res_nl_clipped = fitter_nl_clipped.run()
    i_nl_clipped = res_nl_clipped.beta[0]
    s_nl_clipped = res_nl_clipped.beta[1]
    full_clipped = (-5-i_nl_clipped)/s_nl_clipped*gain
    if full_clipped > np.max(x[valid1])*gain:
        full_clipped = np.max(x[valid1])*gain

    #plot nonlinearity
    plt.subplot(122)
    plt.scatter(x[valid]/1000, nl[valid],color='grey',s=20)
    plt.scatter(x[valid1]/1000, nl[valid1],color='red',s=20)
#    plt.scatter(x[ptc10*valid1]/1000, nl[ptc10*valid1], color='blue', s=20, label='7-70ke-')
#    plt.scatter(x[ptc20*valid]/1000, nl[ptc20*valid], color='black', s=20, label='14-56ke-')
#    plt.scatter(x[ptc20*valid1]/1000, nl[ptc20*valid1], color='red', s=20, label='14-56ke-')
    xplot = np.linspace(0,full_clipped,101)
    plt.plot(xplot/1000, s_nl*xplot+i_nl, color='orange', linestyle='--', label='linear fit')
    plt.plot(xplot/1000, s_nl_clipped*xplot+i_nl_clipped, color='red', linestyle='--', label='clipped linear fit')
    x20 = x[np.where(ptc20*valid0)[0][-1]]
    x10 = x[np.where(ptc10*valid0)[0][-1]]
    nlptc20 = abs(s_nl*x20+i_nl)
    nlptc10 = abs(s_nl*x10+i_nl)
    nlptc20_clipped = abs(s_nl_clipped*x20+i_nl_clipped)
    nlptc10_clipped = abs(s_nl_clipped*x10+i_nl_clipped)
    plt.axhline(y=-3,linestyle='--',color='blue')
    plt.axhline(y=3,linestyle='--',color='blue')
    plt.axhline(y=-5,linestyle='--',color='red')
    plt.axhline(y=5,linestyle='--',color='red')
    plt.axvline(x=full/gain/1000,linestyle='--',color='blue',label='fullwell='+str(np.round(full/1000,1))+"ke-")
    plt.axvline(x=full_clipped/gain/1000,linestyle='--',color='red',label='fullwell_clipped='+str(np.round(full_clipped/1000,1))+"ke-")
    plt.axvline(x=full_image/gain/1000,linestyle='--',color='black',label='fullwell_image='+str(np.round(full_image/1000,1))+"ke-")
    plt.xlabel("Signal (k ADU)")
    plt.ylabel("non-linearity (%)")
    plt.legend()
    
    #save ptc figure
    imfile = fig_outdir+"/"+tab_file.split('/')[-1].replace('.dat','.jpg')
    plt.savefig(imfile)
    return gain, full_image, full, nlptc20, nlptc10, full_clipped, nlptc20_clipped, nlptc10_clipped  
      
def numtrans(n):
    if n<10:
        return '0'+str(n)
    else:
        return str(n)

if __name__ == '__main__':
    if not os.path.exists(out_dir):
        os.system("mkdir "+out_dir)
    if not os.path.exists(out_dir+"/figure"):
        os.system("mkdir "+out_dir+"/figure")
    if not os.path.exists(out_dir+"/data"):
        os.system("mkdir "+out_dir+"/data")
    gains = []
    full_images = []
    fulls = []
    nlptc20s = []
    nlptc10s = []
    fulls_clipped = []
    nlptc20s_clipped = []
    nlptc10s_clipped = []
    fulls_final = []
    for i in range(16):
        num = numtrans(i+1)
        cat = glob(sys.argv[1]+"/ptc_"+wave+"*-c"+num+".dat")[0]
        gain, full_image, full, nlptc20, nlptc10, full_clipped, nlptc20_clipped, nlptc10_clipped = ptc_fit(cat, out_dir+"/figure")
        gains.append(gain)
        full_images.append(full_image)
        fulls.append(full)
        nlptc20s.append(nlptc20)
        nlptc10s.append(nlptc10)
        fulls_clipped.append(full_clipped)
        nlptc20s_clipped.append(nlptc20_clipped)
        nlptc10s_clipped.append(nlptc10_clipped)
        fulls_final.append(min([full_image,full_clipped]))
    tab = Table()
    tab['gain'] = gains
    tab['gain'].unit = 'e-/ADU'
    tab['fullwell_image'] = full_images
    tab['fullwell_image'].unit = 'e-/ADU'
    tab['fullwell'] = fulls
    tab['fullwell'].unit = 'e-'
    tab['nlptc20'] = nlptc20s
    tab['nlptc20'].unit = '%'
    tab['nlptc10'] = nlptc10s
    tab['nlptc10'].unit = '%'
    tab['fullwell_clipped'] = fulls_clipped
    tab['fullwell_clipped'].unit = 'e-'
    tab['full_clipped'] = fulls
    tab['full_clipped'].unit = 'e-'
    tab['full_final'] = fulls_final
    tab['full_final'].unit = 'e-'
    tab['nlptc20_clipped'] = nlptc20s_clipped
    tab['nlptc20_clipped'].unit = '%'
    tab['nlptc10_clipped'] = nlptc10s_clipped
    tab['nlptc10_clipped'].unit = '%'
    tab.write(out_dir+"/data/"+"ptc_"+wave+".tab",format='ipac',overwrite=True)
